﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Public.Ocean.Petrel.Extensions.Internal
{
    public static class Enumerations
    {
        /// <summary>
        /// Flattens a tree structure.
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="inputs">The inputs roots.</param>
        /// <param name="enumerate">Enumerates the children of an element.</param>
        /// <returns></returns>
        public static IEnumerable<T> Flatten<T>(IEnumerable<T> inputs, Func<T, IEnumerable<T>> enumerate)
        {
            if (inputs != null)
            {
                var stack = new Stack<T>(inputs);
                while (stack.Count > 0)
                {
                    var current = stack.Pop();
                    if (current == null) continue;
                    yield return current;
                    var enumerable = enumerate != null ? enumerate(current) : null;
                    if (enumerable != null)
                    {
                        foreach (var child in enumerable)
                        {
                            stack.Push(child);
                        }
                    }
                }
            }
        }

        /// <summary>
        /// Flattens a tree structure.
        /// </summary>
        /// <param name="inputs">The inputs roots.</param>
        /// <param name="enumerate">Enumerates the children of an element.</param>
        /// <returns></returns>
        public static System.Collections.IEnumerable Flatten(System.Collections.IEnumerable inputs, Func<object, System.Collections.IEnumerable> enumerate)
        {
            return Flatten<object>(inputs.Cast<object>(), o => (enumerate(o) ?? new object[0]).Cast<object>());
        }
    }
}
